package code.gateplugins.libsvm;

import gate.Annotation;
import gate.AnnotationSet;
import gate.Controller;
import gate.FeatureMap;
import gate.Resource;
import gate.creole.AbstractLanguageAnalyser;
import gate.creole.ControllerAwarePR;
import gate.creole.ExecutionException;
import gate.creole.ResourceInstantiationException;
import gate.creole.metadata.CreoleParameter;
import gate.creole.metadata.CreoleResource;
import gate.creole.metadata.Optional;
import gate.creole.metadata.RunTime;
import gate.util.OffsetComparator;

import java.io.OutputStream;
import java.io.PrintStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.TreeMap;

import org.apache.log4j.Logger;

import libsvm.svm;
import libsvm.svm_model;
import libsvm.svm_node;
import libsvm.svm_parameter;
import libsvm.svm_problem;

@CreoleResource(name = "LIBSVM classifier PR")
public class SvmClassifierPr extends AbstractLanguageAnalyser implements ControllerAwarePR {
	private static final long serialVersionUID = 1L;
	private static final Logger LOG = Logger.getLogger(SvmClassifierPr.class);
	private SvmData svmData;
	private svm_parameter params;
	private URL modelPath;
	public static enum MODES {
		TRAIN, EVAL, DUMP,
		/** When the pipeline is continuously executed on sets of docs to dump one large features file 
		 * for evaluation, alphabets should not be reloaded */
		CONTINUOUS_DOC_EVAL_DUMP
	};

	// PR parameters
	private MODES mode;
	private URL modelPathParam;
	private String modelCorpusFeature;
	private String inputAS;
	private String labelFeature;
	private String annType;
	private String consoleParams;
	private boolean quietMode;
	private String scoresFeature;

	@Override
	public Resource init() {
		return this;
	}

	@Override
	public void controllerExecutionStarted(Controller c) throws ExecutionException {
		// set model path
		if(modelPathParam == null) {
			String pathValue = (String) corpus.getFeatures().get(modelCorpusFeature);
			if(pathValue == null || pathValue.isEmpty()) {
				throw new ExecutionException("modelCorpusFeature points to an empty feature of the corpus.");
			}
			try {
				modelPath = new URL(pathValue);
			} catch (MalformedURLException e) {
				throw new ExecutionException("modelCorpusFeature is not a valid url!", e);
			}
		} else {
			modelPath = modelPathParam;
		}
		
		if (mode.equals(MODES.EVAL)) {
			try {
				svmData = ModelIO.load(modelPath, false);
			} catch (ResourceInstantiationException e) {
				throw new ExecutionException("Could not load model from modelPath.", e);
			}
			LOG.info("Loaded svm model.");
			// clear only instances, reuse alphabets
			// svmData.setInstances(new ArrayList<Instance>());
		} else {
			if(mode.equals(MODES.CONTINUOUS_DOC_EVAL_DUMP)) {
				if(svmData == null) {
					svmData = new SvmData();
				} else {
					svmData.setModel(null);
					svmData.setInstances(new ArrayList<Instance>());
				}
				// only reload alphabets if not already loaded
				try {
					if(svmData.getAlphabet() == null || svmData.getAlphabet().isEmpty()) {
						svmData = ModelIO.load(modelPath, true);
						LOG.info("Loaded alphabets.");
					}
				} catch (ResourceInstantiationException e) {
					throw new ExecutionException("Could not load alphabets from modelPath.", e);
				}
			} else {
				// clear all
				svmData = new SvmData();
				if(mode.equals(MODES.TRAIN)) {
					if(quietMode) {
						PrintStream origOut = System.out;
						System.setOut(new VoidSysOut(origOut));
					}
				}
			}
		}
		// set svm params
		params = CommandLineParser.parseCommandLine(consoleParams.split("\\s+"));
	}

	@Override
	public void execute() throws ExecutionException {
		AnnotationSet inAS = document.getAnnotations(inputAS);
		List<Annotation> sortedAnns = new ArrayList<Annotation>(inAS.get(annType));
		Collections.sort(sortedAnns, new OffsetComparator());

		for (Annotation ann : sortedAnns) {
			FeatureMap feats = ann.getFeatures();
			Object label = feats.remove(labelFeature);
			feats.remove(scoresFeature);
			if (mode.equals(MODES.EVAL)) {
				svm_node[] x = svmData.toNodes(feats);
				Map<String, Integer> labelMap = svmData.getLabelAlphabet();

				if(params.svm_type == svm_parameter.EPSILON_SVR ||
						params.svm_type == svm_parameter.NU_SVR) {
					double y = svm.svm_predict(svmData.getModel(), x);
					feats.put(labelFeature, y);
				} else {
					double[] scores = new double[(labelMap.size() * (labelMap.size() - 1)) / 2];
					double y = svm.svm_predict_values(svmData.getModel(), x, scores);
					//System.out.println(scores);
					setLabelAndScores(feats, labelMap, scores, y);
				}
			} else if(mode.equals(MODES.CONTINUOUS_DOC_EVAL_DUMP)) {
				// don't grow alphabets
				svmData.addInstance(label, feats, false);
			} else {
				// for train or dump grow alphabets
				if (label != null) {
					// real-valued target class y
					svmData.addInstance(label, feats, true);
				}
			}
		}
	}

	private void setLabelAndScores(FeatureMap feats, Map<String, Integer> labelMap, double[] scores, double y) {
		// TODO: make this smarter:
		TreeMap<Integer, String> invertedLabelMap = new TreeMap<Integer, String>();
		for (Entry<String, Integer> e : labelMap.entrySet()) {
			invertedLabelMap.put(e.getValue(), e.getKey());
		}
		feats.put(labelFeature, invertedLabelMap.get((int) y));

		// TODO: why is this not working:
	/*	StringBuilder sb = new StringBuilder();
		for (Entry<Integer, String> e : invertedLabelMap.entrySet()) {
			sb.append(e.getValue()).append(": ").append(scores[e.getKey() - 1]).append(", ");
		}
		feats.put(scoresFeature, sb.toString()); */	
	}

	/** Training happens here. */
	@Override
	public void controllerExecutionFinished(Controller controller) throws ExecutionException {
		if (!mode.equals(MODES.EVAL)) {
			svm_problem problem = svmData.toSvmProblem();
			if (mode.equals(MODES.TRAIN)) {
				long start = System.currentTimeMillis();
				svm_model mod = svm.svm_train(problem, params);
				svmData.setModel(mod);
				long end = System.currentTimeMillis();
				LOG.info("Training took: " + (end - start));

				// save model and alphabets
				ModelIO.save(modelPath, svmData);
			} else if (mode.equals(MODES.DUMP)) {
				ModelIO.dump(modelPath, problem, svmData, false);
			} else if (mode.equals(MODES.CONTINUOUS_DOC_EVAL_DUMP)) {
				ModelIO.dump(modelPath, problem, svmData, true);
			}
		}
	}

	@RunTime
	@CreoleParameter(comment = "Mode", defaultValue = "TRAIN")
	public void setMode(MODES mode) {
		this.mode = mode;
	}

	@RunTime
	@CreoleParameter(comment = "File path to load/save model or dump libsvm data. " +
			"Model and alphabets are saved in separate files.",
			disjunction = "modelPath")
	public void setModelPath(URL modelPathParam) {
		this.modelPathParam = modelPathParam;
	}
	
	@RunTime
	@CreoleParameter(comment = "Alternative to modelPath, storing the path in corpus feature with this name.",
			disjunction = "modelPath")
	public void setModelCorpusFeature(String modelCorpusFeature) {
		this.modelCorpusFeature = modelCorpusFeature;
	}

	@Optional
	@RunTime
	@CreoleParameter(comment = "Name of annotationSet used for input.")
	public void setInputAS(String inputAS) {
		this.inputAS = inputAS;
	}

	@RunTime
	@CreoleParameter(comment = "Name of the feature which holds the label/class.", defaultValue = "label")
	public void setLabelFeature(String labelFeature) {
		this.labelFeature = labelFeature;
	}

	@RunTime
	@CreoleParameter(comment = "Annotation to be classified.", defaultValue = "Sentence")
	public void setAnnType(String annType) {
		this.annType = annType;
	}
	
	@RunTime
	@CreoleParameter(comment = "Parameters used by libsvm. The string is parsed the same way " +
			"libsvm parses command line options (without the files and quiet mode).", 
			defaultValue = "-s 0 -c 1 -t 0 -m 1024")
	public void setCommandLineParams(String commandLineParams) {
		this.consoleParams = commandLineParams;
	}

	@RunTime
	@CreoleParameter(comment = "Disables libsvm console output.", 
			defaultValue = "false")
	public void setQuietMode(Boolean quietMode) {
		this.quietMode = quietMode;
	}
	
	@RunTime
	@CreoleParameter(comment = "Name of feature used to write result scores.", 
			defaultValue = "scores")
	public void setScoresFeature(String scoresFeat) {
		this.scoresFeature = scoresFeat;
	}

	@Override
	public void controllerExecutionAborted(Controller arg0, Throwable arg1) throws ExecutionException {
		// do nothing
	}

	@Override
	public void reInit() throws ResourceInstantiationException {
		init();
	}

	public String getLabelFeature() {
		return labelFeature;
	}

	public URL getModelPath() {
		return modelPathParam;
	}

	public String getModelCorpusFeature() {
		return modelCorpusFeature;
	}

	
	public MODES getMode() {
		return mode;
	}

	public String getInputAS() {
		return inputAS;
	}

	public String getAnnType() {
		return annType;
	}

	public String getCommandLineParams() {
		return consoleParams;
	}

	public Boolean getQuietMode() {
		return quietMode;
	}

	private class VoidSysOut extends PrintStream {
		public VoidSysOut(OutputStream out) {
			super(out, true);
		}
		@Override
		public void print(String s) {
			// do nothing
		}
	}

	public String getScoresFeature() {
		return scoresFeature;
	}
}